package com.whcoding.base.project.common;

import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.annotation.Resource;
import java.util.List;
import java.util.function.BiFunction;
import java.util.function.ToIntFunction;

/**
 * @program: spring-boot-learning
 * @description: Mybatis 批量插入: https://mp.weixin.qq.com/s/9x2GgAIfeo-y4xKFBsckqQ
 * @author: whcoding
 * @create: 2022-05-05 14:10
 **/
@Component
public class MybatisBatchUtils {

	private static final Logger log = LoggerFactory.getLogger(MybatisBatchUtils.class);

	/**
	 * 每次处理1000条
	 */
	private static final int BATCH_SIZE = 1000;

	@Resource
	private SqlSessionFactory sqlSessionFactory;

	/**
	 * 推荐使用：效率更高
	 * 批量处理修改或者插入
	 * 我们知道上面我们提到了BatchExecutor执行器，
	 * 我们知道每个SqlSession都会拥有一个Executor对象，这个对象才是执行 SQL 语句的幕后黑手，
	 * 我们也知道Spring跟Mybatis整合的时候使用的SqlSession是SqlSessionTemplate，
	 * 默认用的是ExecutorType.SIMPLE，这个时候你通过自动注入获得的Mapper对象其实是没有开启批处理的
	 *
	 * @param data        需要被处理的数据
	 * @param mapperClass Mybatis的Mapper类
	 * @param function    自定义处理逻辑
	 * @return int 影响的总行数
	 */
	public <T, U, R> int batchUpdateOrInsert(List<T> data, Class<U> mapperClass, BiFunction<T, U, R> function) {
		int i = 1;
		SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
		try {
			//获取mapper注入
			U mapper = batchSqlSession.getMapper(mapperClass);
			int size = data.size();
			for (T element : data) {
				function.apply(element, mapper);
				if ((i % BATCH_SIZE == 0) || i == size) {
					//刷新执行结果
					batchSqlSession.flushStatements();
				}
				i++;
			}
			// 非事务环境下强制commit，事务情况下该commit相当于无效
			batchSqlSession.commit(!TransactionSynchronizationManager.isSynchronizationActive());
		} catch (Exception e) {
			batchSqlSession.rollback();
			throw new RuntimeException(e);
		} finally {
			//关闭连接
			batchSqlSession.close();
		}
		return i - 1;
	}


	/**
	 * 不推荐使用
	 * 批量处理修改或者插入
	 *
	 * @param data     需要被处理的数据
	 * @param function 自定义处理逻辑
	 * @return int 影响的总行数
	 */
	public <T> int batchUpdateOrInsert(List<T> data, ToIntFunction<T> function) {
		int count = 0;
		//开启连接
		SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
		try {
			for (int index = 0; index < data.size(); index++) {
				count += function.applyAsInt(data.get(index));
				if (index != 0 && index % BATCH_SIZE == 0) {
					batchSqlSession.flushStatements();
				}
			}
			//提交事务
			batchSqlSession.commit();
		} catch (Exception e) {
			//事务回滚
			batchSqlSession.rollback();
			log.error(e.getMessage(), e);
		} finally {
			//关闭连接
			batchSqlSession.close();
		}
		return count;
	}

}
